extern crate proc_macro;
use ansi_term::Colour;
use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, ItemFn, LitStr};

#[proc_macro_attribute]
pub fn color(args: TokenStream, item: TokenStream) -> TokenStream {
    let mut input = parse_macro_input!(item as ItemFn);
    let fn_name = &input.sig.ident;
    let original_block = &input.block;

    let mut color_name:&str = if !args.is_empty() {
        let arg = parse_macro_input!(args as LitStr);
        &*arg.value()
    } else {
        "Blue" // 默认颜色为蓝色
    };
    let color_enum = match color_name {
        "Red" | "red" => Colour::Red,
        "Green" | "green" => Colour::Green,
        "Blue" | "blue" => Colour::Blue,
        "Yellow" | "yellow" => Colour::Yellow,
        _ => panic!("Unsupported color: {}", color_name),
    };
    let color_enum_path = match color_enum {
        Colour::Red => quote! { crate::Color::Red },
        Colour::Green => quote! { crate::Color::Green },
        Colour::Blue => quote! { crate::Color::Blue },
        Colour::Yellow => quote! { crate::Color::Yellow },
        Colour::Black => quote! { crate::Color::Black },
        Colour::Purple => quote! { crate::Color::Purple },
        Colour::Cyan => quote! { crate::Color::Cyan },
        Colour::White => quote! { crate::Color::White },
        Colour::Fixed(_) => panic!("Fixed color is not supported"),
        Colour::RGB(_, _, _) => panic!("RGB color is not supported"),
    };

    let colored_output = quote! {
        {
            println!("{}", #color_enum_path.paint("! Executing function: ".to_string() + stringify!(#fn_name)).to_string().to_string);
            #original_block
        }
    };

    input.block = syn::parse2(colored_output).unwrap();
    TokenStream::from(quote! { #input })
}
